from Voicelab.pipeline.Node import Node
import parselmouth
from parselmouth.praat import call
from Voicelab.toolkits.Voicelab.VoicelabNode import VoicelabNode
from Voicelab.toolkits.Voicelab.MeasureFormantNode import MeasureFormantNode
from Voicelab.VoicelabWizard.VoicelabTab import VoicelabTab
from Voicelab.VoicelabWizard.F1F2PlotWindow import F1F2PlotWindow

from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *


import io
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.patches import Ellipse
import matplotlib.transforms as transforms

from numpy import random
from scipy.spatial import distance

import numpy as np
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from scipy import stats


###################################################################################################
# F1F2PlotNode
# WARIO pipeline node for estimating the vocal tract of a voice.
###################################################################################################
# ARGUMENTS
# 'voice'   : sound file generated by parselmouth praat
# 'state'   : saves formant data for each voice for processing in end method
###################################################################################################
# RETURNS   : an unnamed string to keep the pipeline running
#           : saves an image file to disk: 'f1f2.png
###################################################################################################


class F1F2PlotNode(VoicelabNode):
    def __init__(self, *args, **kwargs):
        """
        Args:
            *args:
            **kwargs:
        """
        super().__init__(*args, **kwargs)

        self.args = {
            "Vowel Marker": ("Vowels",
                             ["Vowels",
                              "IPA Symbols",
                              "None",
                              ]),
            "Comparison Data": ("Peterson & Barney 1952",
                                ["Peterson & Barney 1952", ]),
        }

        self.state = {
            "f1 means": [],
            "f2 means": [],
            "f3 means": [],
            "f4 means": [],
        }


    def process(self):
        f1 = self.args["F1 Mean"]
        f2 = self.args["F2 Mean"]
        f3 = self.args["F3 Mean"]
        f4 = self.args["F4 Mean"]


        self.state["f1 means"].append(f1)
        self.state["f2 means"].append(f2)
        self.state["f3 means"].append(f3)
        self.state["f4 means"].append(f4)

        f1_means = self.state["f1 means"]
        f2_means = self.state["f2 means"]
        f3_means = self.state["f3 means"]
        f4_means = self.state["f4 means"]

        return {
        }


    def closest_node(self, point, points):
        closest_index = distance.cdist([point], points).argmin()
        return points[closest_index], closest_index


    def end(self, results):
        vowel_marker = self.args["Vowel Marker"][0]
        # create the plot
        fig, ax_kwargs = plt.subplots()
        ax_kwargs.axvline(c='grey', lw=1)
        ax_kwargs.axhline(c='grey', lw=1)

        # Get the Peterson-Barney data
        df = gather_data()

        # Make a list of vowels
        vowels = list(set(df["Vowel"]))

        # Use the list of vowels to specify a colour map for ellipses
        colours = cm.plasma(np.linspace(0, 1, len(vowels)))
        ells = []

        # Convert Our formant measurements to bark
        f1_bark = [hz_to_bark(f1) for f1 in self.state['f1 means']]
        f2_bark = [hz_to_bark(f2) for f2 in self.state['f2 means']]

        # convert peterson barney f1 f2 data into numpy arrays
        search_data = df[["F1 Bark", "F2 Bark"]].values
        #ipa_conversion_dataframe = pd.read_csv('Voicelab/toolkits/Voicelab/IPA_Praat_symbols.csv', header=0)
        ipa_conversion_dataframe = gather_data()
        for point in zip(f1_bark, f2_bark):
            # Find the closest vowel to our data in peterson barney
            node, indexer = (self.closest_node(point, search_data))
            if vowel_marker == "Vowels":
                marker = f"{(df['Vowel'].loc[[indexer]]).values[0]}"
            elif "I" in vowel_marker:
                ipa_symbol_location = \
                    ipa_conversion_dataframe[ipa_conversion_dataframe["Praat"] == (df["IPA"].loc[[indexer]]).values[0]]
                try:
                    ipa_symbol = ipa_symbol_location["IPA"].values.flat[0]
                except:
                    ipa_symbol = '*'
                marker = f"{ipa_symbol}"
            else:
                marker = '.'

            y, x = point
            ax_kwargs.text(x, y, marker, fontsize=12)

        plot_points = pd.DataFrame(list(zip(f2_bark, f1_bark)))
        ax_kwargs.plot(plot_points, marker='o', c='w', alpha=0, linestyle="None")

        for i, vowel in enumerate(vowels):
            y, x = df['F1 Bark'][df['Vowel'] == vowel], df['F2 Bark'][df['Vowel'] == vowel]
            self.confidence_ellipse(x, y, ax_kwargs, n_std=2, alpha=0.3, facecolor=colours[i],
                                    edgecolor='black', zorder=0)
            ells.append(self.ellipse)


        ax_kwargs.legend(ells, vowels)
        ax_kwargs.set_title(f'F1 F2 Plot')

        plt.xlim(left=4)
        plt.xlim(right=20)
        plt.ylim(bottom=2)
        plt.ylim(top=10)
        ax_kwargs.set_xlabel("F2 Bark")
        ax_kwargs.set_ylabel("F1 Bark")
        fig.subplots_adjust(hspace=0.25)
        fig1 = plt.gcf()
        plt.show()

        fig1.savefig('f1f2.png', dpi=1000)
        results[i][self]["F1F2 Plot"] = ["Figure Created"]

        self.f1f2plotwindow = F1F2PlotWindow()
        return results


    def confidence_ellipse(self, x, y, ax, n_std=1, facecolor='none', **kwargs):
        """
        Create a plot of the covariance confidence ellipse of *x* and *y*.

        Parameters
        ----------
        x, y : array-like, shape (n, )
            Input data.

        ax : matplotlib.axes.Axes
            The axes object to draw the ellipse into.

        n_std : float
            The number of standard deviations to determine the ellipse's radiuses.

        Returns
        -------
        matplotlib.patches.Ellipse

        Other parameters
        ----------------
        kwargs : `~matplotlib.patches.Patch` properties
        """
        if x.size != y.size:
            raise ValueError("x and y must be the same size")

        cov = np.cov(x, y)
        pearson = cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])
        # Using a special case to obtain the eigenvalues of this
        # two-dimensionl dataset.
        ell_radius_x = np.sqrt(1 + pearson)
        ell_radius_y = np.sqrt(1 - pearson)
        ellipse = Ellipse((0, 0),
                          width=ell_radius_x * 2,
                          height=ell_radius_y * 2,
                          facecolor=facecolor,
                          **kwargs)
        self.g_ell_center = ellipse.get_center()
        self.g_ell_width = ell_radius_x * 2
        self.g_ell_height = ell_radius_y * 2
        self.angle = 45

        # Calculating the stdandard deviation of x from
        # the squareroot of the variance and multiplying
        # with the given number of standard deviations.
        scale_x = np.sqrt(cov[0, 0]) * n_std
        mean_x = np.mean(x)

        # calculating the stdandard deviation of y ...
        scale_y = np.sqrt(cov[1, 1]) * n_std
        mean_y = np.mean(y)

        transf = transforms.Affine2D() \
            .rotate_deg(45) \
            .scale(scale_x, scale_y) \
            .translate(mean_x, mean_y)

        ellipse.set_transform(transf + ax.transData)
        self.ellipse = ellipse
        return ax.add_patch(ellipse)


def hz_to_bark(hz):
    """
    This function converts Hz to Bark.
        Parameters
        ----------
        hz is the frequency in Hz

        Returns
        -------
        bark is the frequency in bark
    """
    bark = 7 * np.log(hz/650 + np.sqrt(1 + (hz/650)**2))
    return bark


def gather_data():
    """
    This function collects data from Peterson & Barney 1952 from Praat there is no input for the function.

        Returns
        -------
        peterson_barney a pandas dataframe that also includes bark measures using hz_to_bark function
    """

    peterson_barney = call("Create formant table (Peterson & Barney 1952)")
    peterson_barney = pd.read_csv(io.StringIO(call(peterson_barney, "List", True)), sep='\t').dropna()
    peterson_barney['F1 Bark'] = hz_to_bark(peterson_barney['F1'])
    peterson_barney['F2 Bark'] = hz_to_bark(peterson_barney['F2'])
    return peterson_barney
